/*
 * SumLattice.cpp
 *
 *  Created on: Sep 9, 2012
 *      Author: Wanderer
 */

#include "SumLattice.h"

SumLattice::SumLattice(CRF* model, quiver_t* input, value_t* constrain) {
	this->input = input;
	this->constrain = constrain;
	this->length = input->size();
	printf("length: %d\n", length);
	int numStates = model->getNumStates();
	this->pars = model->getParameters();
	this->inipars = model->getInitialParameters();
	this->finpars = model->getFinalParameters();
	this->alpha = new double*[length];
	this->beta = new double*[length];
	this->weights = new double**[length];
	for (int ip = 0; ip < length; ++ip) {
		if (ip != 0) {
			weights[ip] = new double*[numStates];
			for (int i = 0; i < numStates; ++i) {
				weights[ip][i] = new double[numStates];
				for (int j = 0; j < numStates; ++j) {
					weights[ip][i][j] = this->getWeight(ip, j, j);
				}
			}
		}
		alpha[ip] = new double[numStates];
		beta[ip] = new double[numStates];
	}

	//initialize forward variables
	for (int i = 0; i < numStates; ++i) {
		alpha[0][i] = (checkPath(0, i)) ? getInitialWeight(i) : NEG_INF;
		beta[length - 1][i] = (checkPath(length - 1, i)) ? getFinalWeight(i) : NEG_INF;
	}

	//forward
	for (int ip = 1; ip < length; ++ip) {
		for (int i = 0; i < numStates; ++i) {
			alpha[ip][i] = NEG_INF;
			for (int j = 0; j < numStates; ++j) {
				if (checkPath(ip, j)) {
					alpha[ip][i] = sumLog(alpha[ip][i], weights[ip][i][j] + alpha[ip - 1][j]);
				}
			}
		}
	}

	this->totalWeight = NEG_INF;
	for (int i = 0; i < numStates; ++i) {
		totalWeight = sumLog(totalWeight, alpha[length - 1][i]);
	}

	//backward
	for (int i = 0; i < numStates; ++i) {
		gammas[length - 1][i] = alpha[length - 1][i] + beta[length - 1][i] - totalWeight;
	}

	for (int ip = length - 2; ip > 0; ++ip) {
		for (int j = 0; j < numStates; ++j) {
			beta[ip][j] = NEG_INF;
			for (int i = 0; i < numStates; ++i) {
				beta[ip][j] = sumLog(beta[ip][j], weights[ip + 1][i][j] + beta[ip + 1][i]);
			}
			gammas[ip][j] = alpha[ip][j] + beta[ip][j] - totalWeight;
		}
	}

	//free memory
	for (int ip = 0; ip < length; ++ip) {
		if (ip != 0) {
			for (int i = 0; i < numStates; ++i) {
				delete weights[ip][i];
			}
			delete weights[ip];
		}
		delete alpha[ip];
		delete beta[ip];
	}
	delete alpha;
	delete beta;
	delete weights;
}

bool SumLattice::checkPath(int ip, int state) {
	if ((this->constrain == NULL) || (this->constrain->at(ip) == -1) || (state == this->constrain->at(ip))) {
		return true;
	}
	return false;
}

double SumLattice::getWeight(int ip, int i, int j) {
	double weight = 0;
	if (ip == 0) {
		return getInitialWeight(i);
	} else {
		for (unsigned k = 0; k < input->at(ip)->size(); ++k) {
			weight += pars[i][j].at(input->at(ip)->at(k)).value;
		}
	}
	return weight;
}

double SumLattice::getInitialWeight(int i) {
	double weight = 0;
	for (unsigned k = 0; k < input->at(0)->size(); ++k) {
		weight += inipars->at(input->at(0)->at(k)).value;
	}
	return weight;
}

double SumLattice::getFinalWeight(int i) {
	double weight = 0;
	for (unsigned k = 0; k < input->at(length-1)->size(); ++k) {
		if (finpars == NULL)
			return 0;
		weight += finpars->at(input->at(length-1)->at(k)).value;
	}
	return weight;
}

SumLattice::~SumLattice() {
// TODO Auto-generated destructor stub
}

